import pymysql
from tool import logger
import xml.etree.ElementTree as ET
from urllib.parse import urlparse, parse_qs
import os
from config import test_config
import sqlparse

# 获取数据库信息
def get_db_info():
    path = os.path.join(test_config.config["path"]["APP_DIR_PATH"],test_config.config["path"]["APP_MYSQL_INFO_PATH"])
    tree = ET.parse(path)
    root = tree.getroot()
    for item in root.findall('env'):
        name = item.get('name')
        if name == "MYSQL_STRING":
            value = item.get('value')
            return parse_connection_string(value)
    raise Exception("获取数据库信息错误")

# 解析数据库信息
def parse_connection_string(connection_string):
    # 使用 urlparse 解析连接字符串
    parsed = urlparse(connection_string)
    
    # 提取各部分信息
    scheme = parsed.scheme.split('+')[0]  # 数据库类型
    driver = parsed.scheme.split('+')[1]  # 驱动
    username = parsed.username  # 用户名
    password = parsed.password  # 密码
    host = parsed.hostname  # 主机
    port = parsed.port  # 端口
    database = parsed.path.lstrip('/')  # 数据库名
    query_params = parse_qs(parsed.query)  # 查询参数（如字符集）
    
    # 返回解析后的信息
    return {
        "scheme": scheme,
        "driver": driver,
        "user": username,
        "password": password,
        "host": host,
        "port": port,
        "database": database,
        "query_params": {k: v[0] for k, v in query_params.items()}  # 转换为简单字典
    }

# 备份数据库
def backup_database(output_file):
    try:
        mysql_info = get_db_info()
        # 连接到数据库
        connection = pymysql.connect(
            host=mysql_info["host"],
            port=mysql_info["port"],
            user=mysql_info["user"],
            password=mysql_info["password"],
            database=mysql_info["database"],
        )
        cursor = connection.cursor()

        # 打开备份文件
        with open(output_file, "w", encoding="utf-8") as f:
            # 获取所有表名
            cursor.execute("SHOW TABLES;")
            # 检查electrodemapping表是否存在
            tables = [row[0] for row in cursor.fetchall()]
            for table in tables:
                # 导出表结构
                cursor.execute(f"SHOW CREATE TABLE {table};")
                create_table_sql = cursor.fetchone()[1]
                f.write(f"{create_table_sql};\n\n")

                # 获取字段信息，用于识别 Generated Column
                # SHOW FULL COLUMNS 返回的每一行形如:
                # (Field, Type, Collation, Null, Key, Default, Extra, Privileges, Comment)
                cursor.execute(f"SHOW FULL COLUMNS FROM `{table}`;")
                columns_info = cursor.fetchall()

                # 找出所有 Generated 列的名称（有些是 VIRTUAL GENERATED, 有些是 STORED GENERATED）
                # 可以通过判断 columns_info 中的 Extra 字段里是否包含 "GENERATED"
                generated_columns = set(
                    row[0] for row in columns_info if "GENERATED" in row[6].upper()
                )

                # 此时 all_columns 就是这个表中的所有列名
                all_columns = [row[0] for row in columns_info]

                # 过滤掉 Generated 列
                non_generated_columns = [col for col in all_columns if col not in generated_columns]
                if not non_generated_columns:
                    # 如果所有列都是 Generated，那就没法导数据，直接写个空行或者跳过
                    f.write(f"-- Table {table} has no non-generated columns to dump.\n\n")
                    continue

                # 拼接要查询的列名
                select_cols = ", ".join(f"`{col}`" for col in non_generated_columns)
                cursor.execute(f"SELECT {select_cols} FROM `{table}`;")
                rows = cursor.fetchall()
                if rows:
                    columns = ["`{}`".format(desc[0]) for desc in cursor.description]
                    for row in rows:
                        values = ', '.join(
                            "'{}'".format(str(value).replace("\\", "\\\\").replace("'", "\\'").replace('\r', '\\r').replace('\n', '\\n')) if value is not None else "NULL"
                            for value in row
                        )
                        f.write(f"INSERT INTO {table} ({', '.join(columns)}) VALUES ({values});\n")
                f.write("\n")
        print(f"数据库备份成功，文件保存为: {output_file}")
    except Exception as e:
        print(f"备份失败: {e}")
    finally:
        if connection:
            connection.close()

# 导入数据库
def import_backup(sql_file):
    try:
        mysql_info = get_db_info()
        # 连接到数据库
        connection = pymysql.connect(
            host=mysql_info["host"],
            port=mysql_info["port"],
            user=mysql_info["user"],
            password=mysql_info["password"],
            database=mysql_info["database"],
        )
        cursor = connection.cursor()

        # 在执行SQL前，可以先创建被引用的表
        with open(sql_file, "r", encoding="utf-8") as f:
            sql_content = f.read()
            sql_statements = sqlparse.split(sql_content)
            
            # 第一次：执行所有不包含外键的CREATE TABLE语句
            for statement in sql_statements:
                statement = statement.strip()
                if statement:
                    try:
                        cursor.execute(statement)
                    except Exception as e:
                        logger.logger.error(f"创建表失败: {statement[:50]}... 错误: {e}")
            
            # 第二次：执行包含外键的CREATE TABLE语句
            for statement in sql_statements:
                statement = statement.strip()
                if statement and 'CREATE TABLE' in statement.upper() and 'FOREIGN KEY' in statement.upper():
                    try:
                        cursor.execute(statement)
                    except Exception as e:
                        logger.logger.error(f"创建带外键的表失败: {statement[:50]}... 错误: {e}")
            

        # 提交事务
        connection.commit()
                
    except Exception as e:
        logger.logger.error(f"导入失败: {e}")
        raise Exception("导入数据库失败")
    finally:
        if connection:
            connection.close()

# 清空数据库
def clear_database():
    try:
        mysql_info = get_db_info()
        # 连接到数据库
        connection = pymysql.connect(
            host=mysql_info["host"],
            port=mysql_info["port"],
            user=mysql_info["user"],
            password=mysql_info["password"],
            database=mysql_info["database"],
        )
        cursor = connection.cursor()

        # 禁用外键约束
        cursor.execute("SET FOREIGN_KEY_CHECKS = 0;")

        # 获取所有表名
        cursor.execute("SHOW TABLES")
        tables = cursor.fetchall()

        # 删除所有表
        for (table_name,) in tables:
            try:
                cursor.execute(f"DROP TABLE IF EXISTS `{table_name}`")  # 删除表
            except Exception as e:
                logger.logger.error(f"删除表失败: {table_name} 错误: {e}")

        # 启用外键约束
        cursor.execute("SET FOREIGN_KEY_CHECKS = 1;")
        
        connection.commit()
    except Exception as e:
        logger.logger.error(f"清空数据库失败: {e}")
        raise Exception("清空数据库失败")
    finally:
        if connection:
            connection.close()


